// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _EQUATION_SET_HPP
#define _EQUATION_SET_HPP

#include "Panzer_EquationSet_DefaultImpl_decl.hpp"
#include "MaterialEvaluator.hpp"

// include evaluators here
#include "Panzer_Workset_Utilities.hpp"
#include "Panzer_HierarchicParallelism.hpp"
#include "Panzer_STK_GatherFields.hpp"

namespace SiYuan {

  template <typename EvalT>
  class EquationSet : public panzer::EquationSet_DefaultImpl<EvalT>
  {
    public:
      EquationSet(const Teuchos::RCP<Teuchos::ParameterList>& params,
		       const int& default_integration_order,
		       const panzer::CellData& cell_data,
		       const Teuchos::RCP<panzer::GlobalData>& gd,
		       const bool build_transient_support)
      : panzer::EquationSet_DefaultImpl<EvalT>(params,default_integration_order,cell_data,gd,build_transient_support ), m_gd(gd)
      {};

      // ***********************************************************************
      void buildAndRegisterClosureModelEvaluators(PHX::FieldManager<panzer::Traits>& fm,
                                       const panzer::FieldLayoutLibrary& fl,
                                       const Teuchos::RCP<panzer::IntegrationRule>& ir,
                                       const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& factory,
                                       const std::string& model_name,
                                       const Teuchos::ParameterList& models,
                                       const Teuchos::ParameterList& user_data) const
      {
        Teuchos::RCP<Teuchos::ParameterList> pl = this->getEvaluatorParameterList();
        Teuchos::RCP< std::vector< Teuchos::RCP<PHX::Evaluator<panzer::Traits> > > > evaluators = 
          factory.getAsObject<EvalT>()->buildClosureModels(model_name,
                                                     models,
                                                     fl,
                                                     ir,
                                                     *pl,
                                                     user_data,
                                                     m_gd,
                                                     fm);

        auto mevals = buildMaterialEvaluators(ir);
        for( auto a: mevals) {
          evaluators->emplace_back(a);
        }
  
        for (std::vector< Teuchos::RCP<PHX::Evaluator<panzer::Traits> > >::size_type i=0; i < evaluators->size(); ++i)
          this->template registerEvaluator<EvalT>(fm, (*evaluators)[i]);
      }

      /** Alert the panzer library of a DOF provided by this equation set.
      * This automatically sets up the gather/scatter routines neccessary
      * to evaluate and assemble with this unknown.
      *
      * \param[in] dofName (Required) Name of field to lookup in the unique global
      *                        indexer. This also serves as a key for the remaining
      *                        <code>addDOF*</code> methods.
      * \param[in] basisType (Required) Name of the basis type for this DOF.
      * \param[in] basisOrder (Required) Polynomial order for the basis for this DOF.
      * \param[in] integrationOrder (Optional) Order of the integration rule associated with this
      *                             DOF.  If set to -1 (default), it will use the default
      *                             integration order.
      * \param[in] residualName (Optional) Name of field that is to be scattered associated with
      *                         this DOF.  If not supplied or an empty string used, the
      *                         default is to add the prefix "RESIDUAL_" to the dofName for
      *                         the residual field name.
      * \param[in] scatterName (Optional) Name of the required scatter field associated with
      *                        this DOF.  If not supplied or an empty string used,
      *                        the default is to add the prefix "SCATTER_"
      *                        to the dofName for the scatter field name.
      */
      void addVecDOF(const std::string & dofName,
                const std::string & basisType,
                const int & basisOrder,
                const int integrationOrder = -1,
                const std::string residualName = "",
                const std::string scatterName = "");

      std::vector< Teuchos::RCP<PHX::Evaluator<panzer::Traits> > > 
      buildMaterialEvaluators(const Teuchos::RCP<panzer::IntegrationRule>& ir) const
      {
        std::vector< Teuchos::RCP<PHX::Evaluator<panzer::Traits> > >  evaluators;

      //  const std::map<int,Teuchos::RCP<panzer::IntegrationRule> > & ir_map = this->getIntegrationRules();
      //  panzer::IntegrationRule ir = *((ir_map.begin())->second());

        for( auto a: c0) {
          if( _material->size(a)==0 ) {
            std::cout << "Material Property: " << a << "  NOT FOUND!\n";
            throw std::runtime_error("Material Property not defined");
          }
          Teuchos::RCP< PHX::Evaluator<panzer::Traits> > e =
            Teuchos::rcp( new MaterialEvaluator<EvalT, panzer::Traits>(a, _material,*ir) );
          evaluators.push_back(e);
        }
        if (this->buildTransientSupport()) {
          for( auto a: c2) {
            if( _material->size(a)==0 ) {
              std::cout << "Material Property: " << a << "  NOT FOUND!\n";
              throw std::runtime_error("Material Property not defined");
            }
            Teuchos::RCP< PHX::Evaluator<panzer::Traits> > e =
              Teuchos::rcp( new MaterialEvaluator<EvalT, panzer::Traits>(a, _material,*ir) );
            evaluators.push_back(e);
          }
        }
        if (!c3.empty()) {
          for( auto a: c3) {
            if( _material->size(a)==0 ) {
              std::cout << "Material Property: " << a << "  NOT FOUND!\n";
              throw std::runtime_error("Material Property not defined");
            }
            Teuchos::RCP< PHX::Evaluator<panzer::Traits> > e =
              Teuchos::rcp( new MaterialEvaluator<EvalT, panzer::Traits>(a, _material,*ir) );
            evaluators.push_back(e);
          }
        }

        return evaluators;
      }

      virtual void setMaterial(PHX::FieldManager<panzer::Traits>& fm, 
							std::shared_ptr<SiYuan::Material>& matl) override
      {
        _material = matl;
        //const std::vector<panzer::StrPureBasisPair> & sbNames = itr->getProvidedDOFs();
        // Get a unique list of point rules.  NOTE: This assumes that the
        // same point rules are used for all evaluation types.  In the
        // future we could easily change this by adding another level here
        // to differentiate the point rules for each evaluation type.
        // This would require some refactoring of the physics block
        // registration routines and the workset builder to support all
        // combinations of bases and point rules for each evaluation type.
       /* const std::map<int,Teuchos::RCP<panzer::IntegrationRule> > & ir_map = this->getIntegrationRules();
        panzer::IntegrationRule ir = *((ir_map.begin())->second());

        for( auto a: c0) {
          Teuchos::RCP< PHX::Evaluator<panzer::Traits> > jtc =
            Teuchos::rcp( new MaterialEvaluator<EvalT, panzer::Traits>(a, matl,ir) );
          this->template registerEvaluator<EvalT>(fm, jtc);
        }*/

      }

      std::vector< Teuchos::RCP<PHX::Evaluator<panzer::Traits> > > 
      buildMeshinputFieldEvaluators(PHX::FieldManager<panzer::Traits>& fm,
                                    const panzer_stk::STK_Interface& mesh) const
      {
        std::vector< Teuchos::RCP<PHX::Evaluator<panzer::Traits> > >  evaluators;

        if( _mesh_fields.empty() ) return evaluators;    // no mesh field needed
        auto fdvec = mesh.getMetaData()->get_fields( mesh.getNodeRank() );
        std::size_t n_fields = fdvec.size();
        if( n_fields<=1 ) return evaluators;             // no field defined in mesh

        for( auto s: _mesh_fields ) {
          for( std::size_t i=0; i<n_fields; i++ ) {
            if( fdvec[i]->name() == s ) {
              Teuchos::ParameterList p;
              Teuchos::RCP< PHX::Evaluator<panzer::Traits> > e = 
                Teuchos::rcp( new panzer_stk::GatherFields<EvalT, panzer::Traits>(mesh,p));
              evaluators.push_back(e);
              break;
            }
          }
        }

        return evaluators;
      }

    protected:
      std::vector<std::string> c0,c1,c2,c3;   // coefficients of du^2/dx^2, du/dx, du/dt, du^2/dt^2 if needs. Otherwise all in c0
      std::string model_id;
      std::shared_ptr<Material> _material;
      std::vector<std::string> _mesh_fields;  // fields needed, which may defined by mesh file
	  Teuchos::RCP<panzer::GlobalData> m_gd;

      void addC0(std::string a) {c0.emplace_back(a);}
      void addC1(std::string a) {c1.emplace_back(a);}
      void addC2(std::string a) {c2.emplace_back(a);}
      void addC3(std::string a) {c3.emplace_back(a);}

      std::vector<std::string> getC0() const {return c0;}
      std::vector<std::string> getC1() const {return c1;}
      std::vector<std::string> getC2() const {return c2;}
      std::vector<std::string> getC3() const {return c3;}

      void addMeshField(std::string s) {_mesh_fields.emplace_back(s);}
  };
  
}

#endif
